(* Martin Jambon, Yoann Padioleau
 *
 * Copyright (C) 2024-2025 Semgrep Inc.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * version 2.1 as published by the Free Software Foundation, with the
 * special exception on linking described in file LICENSE.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the file
 * LICENSE for more details.
 *)
open Common

(*****************************************************************************)
(* Prelude *)
(*****************************************************************************)
(* Centralize and unify the use of the standard output (and error) of a
 * program and automatically handle whether or not to display colors
 * depending on the terminal capabilities and/or CLI flags.
 *
 * There are many different ways a program can output a string on
 * the standard output:
 *  - Stdlib.print_string using Stdlib.stdout
 *  - Printf.printf using Stdlib.stdout
 *  - Unix.write using Unix.stdout
 *  - Format.fmt using Formatter.std_formatter
 *  - Fmt.pr, or Fmt.pf using Fmt.stdout
 *  - ANSITerminal.print_string
 *  - Ocolor_format.printf
 *  - UCommon.pr
 *  - ...
 *  - and now UConsole.print() and CapConsole.print too :)
 *
 * (See also TCB/forbid_console.jsonnet for the list of functions we
 * want to forbid to enforce instead the use of this module and CapConsole.ml)
 *
 * One of the goal of this module is to provide another way that can be
 * mocked and grepped and redirected more easily. This is useful
 * for command-line programs that want to clearly separate
 * the normal output of the program from other output.
 *
 * alt:
 *  - Logs.app(), but Logs.app actually print on stderr by default
 *    (without any leading tag), and using a Logs.xxx function for
 *    a regular output feels a bit weird
 *  - use Common.pr(), but we can't easily redirect or mock the output
 *    and we use Common.pr in too many contexts. The use of a separate
 *    UConsole.print() makes it clear this is intended to be the real output
 *    of the program.
 *  - use directly ANSITerminal.ml but again can't be mocked, does not
 *    play well with capabilities so we need to wrap it anyway (CapConsole.ml),
 *    and does not automatically disable colors depending on a highlight config.
 *  - use Fmt but Format pretty printing is too complex for most use cases
 *    and not really needed.
 *
 * TODO:
 *  - use this module for printing all user-facing messages.
 *    Maybe it's only possible once we drop pysemgrep because of bugs and
 *    quirks we have to reproduce for now.
 *  - allow to mock the output, with_mock_output () ?
 *    update: or use Testo recent capture_stdout which uses
 *    some Unix.dupe internally so no need to mock?
 *  - allow to redirect the output via a --output flag
 *    with_redirect_output () ?
 *
 * The 'Logs_' module depends on this module because logging
 * is done on stderr by default.
 *
 * This module contains the shared part between UConsole.ml and
 * CapConsole.ml
 *
 * Some notes on Windows:
 * The default terminal of Windows does not handle ANSI sequences and so
 * what is generated by ANSITerminal.sprintf(). There are ways to print
 * colors on such Windows terminal but the API is annoying (one can't
 * use ANSITerminal.sprintf but have to use ANSITerminal.printf which is less
 * composable) and anyway if one really wants colors on Windows, one can
 * install fancier and more compatible terminal (or use VSCode which integrates
 * a builtin ANSI-compliant terminal).
 *)

(*****************************************************************************)
(* Types and globals *)
(*****************************************************************************)

(* Usually derived from CLI flags (e.g., --force-color=true) *)
type highlight_setting = Auto | On | Off [@@deriving show]

(* The result of applying 'highlight_setting' *)
type highlight = On | Off [@@deriving show]

(* nosemgrep: no-ref-declarations-at-top-scope *)
let highlight_setting : highlight_setting ref = ref Auto

(* nosemgrep: no-ref-declarations-at-top-scope *)
let highlight : highlight ref = ref (Off : highlight)

(* accessors *)
let get_highlight_setting () = !highlight_setting
let get_highlight () = !highlight

let with_highlight_setting ~isatty setting func =
  let orig_setting = !highlight_setting in
  let orig_highlight = !highlight in
  let hl : highlight =
    match setting with
    | Auto -> if isatty then On else Off
    | On -> On
    | Off -> Off
  in
  highlight_setting := setting;
  highlight := hl;
  Common.protect func ~finally:(fun () ->
      highlight_setting := orig_setting;
      highlight := orig_highlight)

(*****************************************************************************)
(* Partial copy of ANSITerminal.ml *)
(*****************************************************************************)
(* So one can use Console.Red instead of ANSITerminal.Red *)

type color = ANSITerminal.color =
  | Black
  | Red
  | Green
  | Yellow
  | Blue
  | Magenta
  | Cyan
  | White
  | Default  (** Default color of the terminal *)

type style = ANSITerminal.style =
  | Reset
  | Bold
  | Underlined
  | Blink
  | Inverse
  | Hidden
  | Foreground of color
  | Background of color

let black = ANSITerminal.black
let red = ANSITerminal.red
let green = ANSITerminal.green
let yellow = ANSITerminal.yellow
let blue = ANSITerminal.blue
let magenta = ANSITerminal.magenta
let cyan = ANSITerminal.cyan
let white = ANSITerminal.white
let default = ANSITerminal.default

let sprintf styles =
  match get_highlight () with
  | Off -> Printf.sprintf
  | On -> ANSITerminal.sprintf styles

(* shortcuts *)
let bold s = sprintf [ Bold ] "%s" s
let underline s = sprintf [ Underlined ] "%s" s
let color c s = sprintf [ c ] "%s" s

(*****************************************************************************)
(* Error/Warning/Success messages *)
(*****************************************************************************)

(* hopefully not too semgrep specific *)
type semgrep_style = Error | Warning | Success

let color_of_style = function
  | Error -> red
  | Warning -> yellow
  | Success -> green

let style_string style str =
  match get_highlight () with
  | On -> sprintf [ color_of_style style ] "%s" str
  | Off -> str

(* exported in Console.mli *)
let error str = style_string Error str
let warning str = style_string Warning str
let success str = style_string Success str

(*****************************************************************************)
(* Helpers for table() below *)
(*****************************************************************************)

(* ex: line 5 --> "|---|" ? *)
let line (width : int) : string =
  String.init (3 * width) (fun i ->
      char_of_int
        (match i mod 3 with
        | 0 -> 0xE2
        | 1 -> 0x94
        | 2 -> 0x80
        | _not_possible -> assert false))

(* val layout_table :
    string * string list -> (string * int list) list -> string list
*)
let layout_table (h1, heading) entries =
  let int_size i =
    let rec dec acc = function
      | 0 -> acc
      | n -> dec (succ acc) (n / 10)
    in
    if i =|= 0 then 1 else dec 0 i
  in
  let len1, lengths =
    let acc = List_.map String.length heading in
    List.fold_left
      (fun (n1, needed) (c1, curr) ->
        ( max (String.length c1) n1,
          List_.map2_exn max needed (List_.map int_size curr) ))
      (String.length h1, acc)
      entries
  in
  let lengths = List_.map (fun i -> i + 3) lengths in
  let line = List.fold_left (fun acc w -> acc + w) (len1 + 2) lengths |> line in
  let pad str_size len =
    let to_pad = len - str_size in
    String.make to_pad ' '
  in
  String.concat ""
    (List_.flatten
       ([ h1; pad (String.length h1) len1 ]
       :: List_.map2_exn
            (fun h l -> [ pad (String.length h) l; h ])
            heading lengths))
  :: line
  :: List_.map
       (fun (e1, entries) ->
         String.concat ""
           (List_.flatten
              ([ e1; pad (String.length e1) len1 ]
              :: List_.map2_exn
                   (fun e l -> [ pad (int_size e) l; string_of_int e ])
                   entries lengths)))
       entries

(*****************************************************************************)
(* Ascii art for headings and tables *)
(*****************************************************************************)

(* old: was Fmt_.pp_heading before but better not using Fmt when not needed *)
let heading (str : string) : string =
  let line = line (String.length str + 2) in
  spf "\n\n┌%s┐\n" line ^ spf "│ %s │\n" str ^ spf "└%s┘\n" line

(* old: was Fmt_.pp_table before *)
let table (h1, heading) entries : string =
  let lines = layout_table (h1, heading) entries in
  Buffer_.with_buffer_to_string (fun buf ->
      let prf fmt = Printf.bprintf buf fmt in
      lines
      |> List.iteri (fun idx line ->
             prf "%s%s\n" (if idx =|= 1 then " " else "  ") line);
      prf "\n")

(* old: was Fmt_.pp_tables before *)
let tables (h1, heading1, entries1) (h2, heading2, entries2) : string =
  let lines1 = layout_table (h1, heading1) entries1
  and lines2 = layout_table (h2, heading2) entries2 in
  let l1_space =
    String.make
      (String.length (List_.hd_exn "unexpected empty list" lines1))
      ' '
  in
  let space = String.make 10 ' ' in
  Buffer_.with_buffer_to_string (fun buf ->
      let prf fmt = Printf.bprintf buf fmt in

      let rec one idx a b =
        match (a, b) with
        | [], [] -> ()
        | a :: ra, b :: rb ->
            prf "%s%s%s%s\n"
              (if idx =|= 1 then " " else "  ")
              a
              (if idx =|= 1 then String.make 8 ' ' else space)
              b;
            one (idx + 1) ra rb
        | a :: ra, [] ->
            prf "  %s\n" a;
            one (idx + 1) ra []
        | [], b :: rb ->
            prf "  %s%s%s\n" l1_space space b;
            one (idx + 1) [] rb
      in
      one 0 lines1 lines2)

(*****************************************************************************)
(* TODO delete at some point *)
(*****************************************************************************)

let strong_style_string style str =
  match get_highlight () with
  | On -> sprintf [ white; Bold; color_of_style style ] "%s" str
  | Off -> str

(* string variants *)
let strong_error str = strong_style_string Error str
let strong_warning str = strong_style_string Warning str
let strong_success str = strong_style_string Success str

(* TODO: get rid of? should use Logs module instead? *)
let error_tag () = strong_error " ERROR "
let warning_tag () = strong_warning " WARNING "
let success_tag () = strong_success " SUCCESS "
